package lib

func init() {
	Probs = append(Probs, Problem{
		Num:         373,
		Discription: "从两个非递减数组里各取一个数字形成一对，求和前k小的对",
		Level:       3,
		Labels: map[string]int{
			"大小堆": 1,
			"多元排序": 1,
			"哈希表": 1,
		},
	})
}

func KSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
	res := make([][]int, 0)
	s := [][]int{{0, 0}}

	isSmaller := func(a []int, b []int) bool {
		return nums1[a[0]]+nums2[a[1]] <= nums1[b[0]]+nums2[b[1]]
	}

	insert := func(val []int) {
		s = append(s, val)
		i := len(s) - 1
		for (i-1)/2 >= 0 {
			parentIndex := (i - 1) / 2
			if isSmaller(s[parentIndex], val) {
				break
			}

			s[i], s[parentIndex] = s[parentIndex], s[i]
			i = parentIndex
		}
	}

	adjust := func() {
        //长度<=1无需调整
		if len(s) <= 1 {
			return
		}

		i := 0
		topVal := s[0]
		n := len(s) - 1
		for 2*i+1 <= n {
			leftIndex := 2*i + 1
			rightIndex := leftIndex + 1
			smallerIndex := leftIndex
			if rightIndex <= n && isSmaller(s[rightIndex], s[leftIndex]) {
				smallerIndex = rightIndex
			}

			if isSmaller(topVal, s[smallerIndex]) {
				break
			}

			s[smallerIndex], s[i] = s[i], s[smallerIndex]
			i = smallerIndex
		}
	}

    //为了避免重复进队
	indexMap := make(map[[2]int]bool, 0)
	for i := 0; i < k && len(s) > 0; i++ {
		a := s[0][0]
		b := s[0][1]
		res = append(res, []int{nums1[a], nums2[b]})
		end := len(s) - 1
		s[0] = s[end]
		s = s[:end]
		adjust()
        //创建新的对时，要注意会不会越原数组的界
		if a <= len(nums1)-2 && !indexMap[[2]int{a + 1, b}] {
			insert([]int{a + 1, b})
			indexMap[[2]int{a + 1, b}] = true
		}
		if b <= len(nums2)-2 && !indexMap[[2]int{a, b + 1}] {
			insert([]int{a, b + 1})
			indexMap[[2]int{a, b + 1}] = true
		}
	}

	return res
}

